
import pandas as pd
import numpy as np
# import matplotlib.pyplot as plt
from util_set_zh_matplot import plt
import seaborn as sns
from pathlib import Path
import pdb
import os

import util_for_output_zh

# 在查看数据前设置
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

# 采用四分位数、五分位数、六分位数三种分法，将BMI划分为多组，确保每组样本量均衡（避免单组样本量＜总样本量的10%）
# PSO 粒子群优化 逐步增加 BMI 阈值 检查 Y浓度>4%时时间点是否更早
import numpy as np
import pandas as pd
from pyswarm import pso  # 粒子群优化库
from sklearn.mixture import GaussianMixture
import matplotlib.pyplot as plt
# 1. 数据加载与预处理 - 优化点：提前筛选和转换
def convert_gestational_week(week_str):
    """快速转换孕周为数值"""
    if 'w+' in week_str:
        w, d = week_str.split('w+')
        return float(w) + float(d)/7
    return float(week_str.split('w')[0])

# 加载数据并立即进行必要筛选
df = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')

# 只处理需要的列，减少数据量
df = df[['孕妇BMI', '检测孕周', 'Y染色体浓度']]

# 提前过滤Y浓度>0.04的数据
df = df[df['Y染色体浓度'] > 0.04].copy()

# 转换孕周并删除原始列
df['检测孕周数值'] = df['检测孕周'].apply(convert_gestational_week)
df = df.drop('检测孕周', axis=1)

# 2. 动态BMI分段（优化版本）
def dynamic_bmi_segmentation(data, n_segments=4, min_ratio=0.1):
    """更高效的BMI分段方法，使用分位数而非GMM"""
    bmi = data['孕妇BMI'].values
    total = len(data)
    min_samples = total * min_ratio
    
    # 先尝试简单分位数
    quantiles = np.linspace(0, 1, n_segments + 1)[1:-1]
    cutoffs = np.percentile(bmi, quantiles * 100)
    
    # 检查是否满足最小样本量要求
    segments = np.digitize(bmi, cutoffs)
    counts = np.bincount(segments, minlength=n_segments)
    
    # 如果有分段样本量不足，调整分界点
    if not all(counts >= min_samples):
        # 手动调整确保每个分段有足够样本
        sorted_bmi = np.sort(bmi)
        cutoffs = []
        for i in range(1, n_segments):
            cutoff_idx = int(min_samples * i)
            cutoffs.append(sorted_bmi[cutoff_idx])
        cutoffs = np.unique(cutoffs)  # 确保唯一性
    
    return np.sort(cutoffs)

# 获取BMI分段并预计算各分段数据 - 优化点：提前计算所有分段数据
n_segments = 4
bmi_cutoffs = dynamic_bmi_segmentation(df, n_segments)
print(f"自动确定的BMI分界点: {bmi_cutoffs}")

# 预计算每个BMI分段的数据 - 关键优化：避免在目标函数中重复计算
seg_data_list = []
for i in range(n_segments):
    if i == 0:
        seg_data = df[df['孕妇BMI'] < bmi_cutoffs[0]]['检测孕周数值']
    elif i == n_segments - 1:
        seg_data = df[df['孕妇BMI'] >= bmi_cutoffs[-1]]['检测孕周数值']
    else:
        seg_data = df[(df['孕妇BMI'] >= bmi_cutoffs[i-1]) & 
                     (df['孕妇BMI'] < bmi_cutoffs[i])]['检测孕周数值']
    seg_data_list.append(seg_data)

# 3. 优化的目标函数 - 优化点：使用预计算数据，减少重复计算
def constrained_objective(x):
    """优化的目标函数，使用预计算的分段数据"""
    # 约束检查 - 放在最前面快速排除无效解
    if not (x[0] <= x[1] <= x[2] <= x[3]):
        return 1e6  # 不可行解
    
    if any(x < 10) or any(x > 30):
        return 1e6
    
    segment_weights = [1.5, 1.25, 1.0, 0.75]
    time_penalty = 0.0
    accuracy_reward = 0.0
    diversity_penalty = 0.0
    
    # 利用预计算的分段数据
    for i in range(n_segments):
        seg_data = seg_data_list[i]
        n = len(seg_data)
        if n == 0:
            continue
            
        # 计算该时间点前的检测率
        acc = np.mean(seg_data <= x[i])
        
        # 时间惩罚（越早越好）
        time_penalty += x[i] * segment_weights[i]
        
        # 准确率奖励（越高越好）
        accuracy_reward -= 100 * acc * segment_weights[i]
    
    # 多样性惩罚（防止分段过于接近）
    for i in range(1, n_segments):
        diff = x[i] - x[i-1]
        if diff < 0.5:  # 最小允许差异
            diversity_penalty += (0.5 - diff) * 100
    
    return time_penalty + accuracy_reward + diversity_penalty

# 4. 运行PSO优化 - 优化点：调整参数减少迭代次数
lb = [10, 12, 14, 16]  # 下限：分段递增
ub = [20, 22, 24, 25]  # 上限：分段递增

# 调整PSO参数以加快收敛
best_x, best_f = pso(
    constrained_objective,
    lb, ub,
    swarmsize=20,       # 减少粒子数量
    maxiter=100,        # 减少迭代次数
    debug=False,        # 关闭调试输出
    phip=0.7, 
    phig=0.8,
    minstep=0.1,        # 增加最小步长
    minfunc=0.1         # 增加最小函数变化
)

# 5. 结果分析与可视化
print("\n=== 优化结果 ===")
for i in range(n_segments):
    lower = bmi_cutoffs[i-1] if i > 0 else -np.inf
    upper = bmi_cutoffs[i] if i < n_segments - 1 else np.inf
    seg_data = seg_data_list[i]
    
    print(f"BMI分段 {lower:.1f}-{upper:.1f}:")
    print(f"  最佳检测时间: {best_x[i]:.2f} 周")
    
    # 计算实际合格率
    acc = np.mean(seg_data <= best_x[i]) if len(seg_data) > 0 else 0
    print(f"  合格率: {acc:.2%}")
    print(f"  该分段样本量: {len(seg_data)}")

# 6. 可视化
plt.figure(figsize=(14, 7))
colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']

for i in range(n_segments):
    # 获取BMI范围
    if i == 0:
        bmi_mask = df['孕妇BMI'] < bmi_cutoffs[0]
    elif i == n_segments - 1:
        bmi_mask = df['孕妇BMI'] >= bmi_cutoffs[-1]
    else:
        bmi_mask = (df['孕妇BMI'] >= bmi_cutoffs[i-1]) & (df['孕妇BMI'] < bmi_cutoffs[i])
    
    # 绘制数据点
    seg_data = df[bmi_mask]
    plt.scatter(
        seg_data['检测孕周数值'], 
        seg_data['孕妇BMI'],
        color=colors[i],
        alpha=0.6,
        label=f'BMI分段{i+1}'
    )
    
    # 绘制最佳时间线
    plt.axvline(
        x=best_x[i],
        color=colors[i],
        linestyle='--',
        linewidth=2,
        label=f'分段{i+1}最佳时间({best_x[i]:.1f}周)'
    )

plt.xlabel('检测孕周', fontsize=12)
plt.ylabel('孕妇BMI', fontsize=12)
plt.title('不同BMI分段的最佳检测时间点优化结果', fontsize=14)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('ques2_optimized_detection_time_with_constraints.png', dpi=300, bbox_inches='tight')
plt.close()